from collections import List

struct DoseStatRet:
    var binned_dose_values: List[Float64]
    var accumulated_relative_volumes: List[Float64]

    fn __init__(inout self):
        self.binned_dose_values = List[Float64]()
        self.accumulated_relative_volumes = List[Float64]()

    fn __copyinit__(inout self, existing: Self):
        self.binned_dose_values = List[Float64](existing.binned_dose_values)
        self.accumulated_relative_volumes = List[Float64](existing.accumulated_relative_volumes)

fn accumulate(
    voxel_volume: List[Float64],
    voxel_dose: List[Float64],
    do_exact_accumulation: Bool = False
) -> DoseStatRet:
    """计算剂量体积直方图(DVH)"""

    var result = DoseStatRet()
    var dose_volume_pairs = List[Tuple[Float64, Float64]]()

    # 过滤体积为0的体素，创建剂量-体积对
    for i in range(len(voxel_volume)):
        if voxel_volume[i] > 0.0:
            var dose = max(0.0, voxel_dose[i])
            dose_volume_pairs.append((dose, voxel_volume[i]))

    var num_positive = len(dose_volume_pairs)

    # 空数据快速处理
    if num_positive == 0:
        result.accumulated_relative_volumes = [1.0, 0.0]
        result.binned_dose_values = [0.0, 0.0]
        return result

    # 算法选择阈值
    var EXACT_THRESHOLD: Int = 5000
    var use_exact_method = do_exact_accumulation or (num_positive < EXACT_THRESHOLD)

    if use_exact_method:
        # 精确计算：直接排序
        sort_dose_volume_pairs(dose_volume_pairs)

        result.binned_dose_values = List[Float64](num_positive + 2)
        result.accumulated_relative_volumes = List[Float64](num_positive + 2)

        # 填充排序后的结果，添加哨兵
        result.binned_dose_values[0] = dose_volume_pairs[0][0]
        result.accumulated_relative_volumes[0] = 0.0

        for i in range(num_positive):
            result.binned_dose_values[i + 1] = dose_volume_pairs[i][0]
            result.accumulated_relative_volumes[i + 1] = dose_volume_pairs[i][1]

        result.binned_dose_values[num_positive + 1] = dose_volume_pairs[num_positive - 1][0]
        result.accumulated_relative_volumes[num_positive + 1] = 0.0

    else:
        # 粗略计算：分箱统计
        var (min_dose, max_dose, total_volume) = calculate_dose_stats(dose_volume_pairs)
        var dose_range = max_dose - min_dose

        # 剂量范围检查
        if dose_range <= 1e-6:
            raise Error("剂量范围过小，请检查输入数据")

        # 动态计算分箱数
        var num_bins = calculate_optimal_bins(dose_range, num_positive)
        var bin_width = dose_range / (num_bins - 1)

        # 初始化结果数组
        result.binned_dose_values = List[Float64](num_bins + 2)
        result.accumulated_relative_volumes = List[Float64](num_bins + 2)

        # 设置剂量分箱值
        for j in range(num_bins):
            result.binned_dose_values[j + 1] = min_dose + j * bin_width

        # 统计每个分箱的体积
        for i in range(num_positive):
            var (dose, volume) = dose_volume_pairs[i]

            if dose <= min_dose:
                result.accumulated_relative_volumes[1] += volume
            elif dose >= max_dose:
                continue  # 忽略最大值哨兵
            else:
                var bin_index = Int((dose - min_dose) / bin_width + 0.5)
                result.accumulated_relative_volumes[bin_index + 1] += volume

        # 设置哨兵值
        result.binned_dose_values[0] = result.binned_dose_values[1]
        result.binned_dose_values[num_bins + 1] = result.binned_dose_values[num_bins]
        result.accumulated_relative_volumes[0] = 0.0
        result.accumulated_relative_volumes[num_bins + 1] = 0.0

    # 计算累积相对体积并归一化
    calculate_cumulative_volumes(inout result)

    return result

fn sort_dose_volume_pairs(inout pairs: List[Tuple[Float64, Float64]]):
    """按剂量值排序剂量-体积对"""
    # 使用简单的插入排序（对于小数据集）或实现更高效的排序
    var n = len(pairs)
    for i in range(1, n):
        var j = i
        while j > 0 and pairs[j][0] < pairs[j-1][0]:
            pairs[j], pairs[j-1] = pairs[j-1], pairs[j]
            j -= 1

fn calculate_dose_stats(pairs: List[Tuple[Float64, Float64]]) -> Tuple[Float64, Float64, Float64]:
    """计算剂量的统计特征"""
    var min_dose = pairs[0][0]
    var max_dose = pairs[0][0]
    var total_volume = 0.0

    for (dose, volume) in pairs:
        min_dose = min(min_dose, dose)
        max_dose = max(max_dose, dose)
        total_volume += volume

    return (max(0.0, min_dose), max_dose, total_volume)

fn calculate_optimal_bins(dose_range: Float64, num_points: Int) -> Int:
    """计算最优分箱数"""
    var MIN_BINS = 2
    var MAX_BINS = 1001
    var dose_min_width = 1e-6

    var theoretical_bins = dose_range / dose_min_width + 1.0
    var num_bins = Int(theoretical_bins)

    num_bins = max(MIN_BINS, min(MAX_BINS, num_bins))
    return num_bins

fn calculate_cumulative_volumes(inout result: DoseStatRet):
    """计算累积体积并归一化"""
    var num_bins = len(result.accumulated_relative_volumes) - 2
    var total_volume = 0.0

    # 计算总体积
    for i in range(1, num_bins + 1):
        total_volume += result.accumulated_relative_volumes[i]

    # 计算累积体积（从高剂量到低剂量）
    var cumulative = 0.0
    for i in range(num_bins, 0, -1):
        var current_volume = result.accumulated_relative_volumes[i]
        result.accumulated_relative_volumes[i] = cumulative + 0.5 * current_volume
        cumulative += current_volume

    # 设置哨兵和归一化
    result.accumulated_relative_volumes[0] = total_volume
    if total_volume > 0.0:
        for i in range(len(result.accumulated_relative_volumes)):
            result.accumulated_relative_volumes[i] /= total_volume

fn get_metric(
    dose_at_volume: Float64,
    binned_dose_values: List[Float64],
    accumulated_relative_volumes: List[Float64]
) -> Float64:
    """获取特定体积百分比对应的剂量"""

    if len(accumulated_relative_volumes) == 0 or len(binned_dose_values) == 0:
        print("警告：输入向量为空")
        return 0.0

    # 边界检查
    if dose_at_volume >= accumulated_relative_volumes[0]:
        return binned_dose_values[0]
    if dose_at_volume <= accumulated_relative_volumes[-1]:
        return binned_dose_values[-1]

    # 查找目标体积对应的位置
    var index = binary_search(accumulated_relative_volumes, dose_at_volume, greater=True)

    if index < len(accumulated_relative_volumes) and accumulated_relative_volumes[index] == dose_at_volume:
        # 精确匹配：计算平均值
        var sum_doses = 0.0
        var count = 0

        var i = index
        while i < len(accumulated_relative_volumes) and accumulated_relative_volumes[i] == dose_at_volume:
            sum_doses += binned_dose_values[i]
            count += 1
            i += 1

        return sum_doses / Float64(count)

    else:
        # 插值计算
        var ind = index
        if ind == 0:
            return binned_dose_values[0]
        if ind >= len(accumulated_relative_volumes):
            return binned_dose_values[-1]

        var ind1 = ind - 1
        var ind2 = ind

        var vol_diff = accumulated_relative_volumes[ind2] - accumulated_relative_volumes[ind1]

        if abs(vol_diff) <= 1e-10:
            return 0.5 * (binned_dose_values[ind2] + binned_dose_values[ind1])

        var dose_diff = binned_dose_values[ind2] - binned_dose_values[ind1]
        return binned_dose_values[ind1] + (dose_at_volume - accumulated_relative_volumes[ind1]) * dose_diff / vol_diff

fn binary_search(
    arr: List[Float64],
    target: Float64,
    greater: Bool = False
) -> Int:
    """二分查找实现"""
    var low = 0
    var high = len(arr) - 1
    var result = len(arr)

    while low <= high:
        var mid = (low + high) // 2

        if greater:
            if arr[mid] > target:
                result = mid
                low = mid + 1
            else:
                high = mid - 1
        else:
            if arr[mid] < target:
                result = mid
                low = mid + 1
            else:
                high = mid - 1

    return result

fn get_dose_stat_roi_interp(
    dose_type: String,
    binned_dose_values: List[Float64],
    accumulated_relative_volumes: List[Float64],
    voxel_volume: List[Float64],
    voxel_dose: List[Float64],
    dose_at_volume: Float64 = 0.0
) -> Float64:
    """计算ROI内的剂量统计指标"""

    if len(voxel_volume) != len(voxel_dose):
        raise Error("体素体积和剂量数组大小不匹配")

    # 剂量类型处理
    if dose_type == "average":
        # 计算平均剂量
        var total_volume = 0.0
        var weighted_dose = 0.0

        for i in range(len(voxel_volume)):
            total_volume += voxel_volume[i]
            weighted_dose += voxel_dose[i] * voxel_volume[i]

        if total_volume <= 0.0:
            raise Error("ROI体积为零，无法计算平均剂量")

        return weighted_dose / total_volume

    # 解析剂量类型
    var target_volume: Float64
    if dose_type == "D_98":
        target_volume = 0.98
    elif dose_type == "D_95":
        target_volume = 0.95
    elif dose_type == "D_50":
        target_volume = 0.5
    elif dose_type == "D_2":
        target_volume = 0.02
    elif dose_type == "DoseAtVolume":
        target_volume = dose_at_volume
    else:
        raise Error("不支持的剂量类型: " + dose_type)

    return get_metric(target_volume, binned_dose_values, accumulated_relative_volumes)